#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/1/23 20:02
# @Author  : Cxk
import h5py
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import numpy as np
from vgg16_keras import VGGNet


def get_img_list(serach_img_path,all_img_number):
    query = serach_img_path #识别图片
    index = 'models/vgg_featureCNN.h5' #文件夹保存的模型
    result = 'img' #分类文件夹

    h5f = h5py.File(index, 'r')
    feats = h5f['dataset_1'][:]
    imgNames = h5f['dataset_2'][:]
    h5f.close()

    model = VGGNet()

    queryVec = model.vgg_extract_feat(query)  # 修改此处改变提取特征的网络
    scores = np.dot(queryVec, feats.T)
    rank_ID = np.argsort(scores)[::-1]
    rank_score = scores[rank_ID]
    # print(rank_score)

    maxres = all_img_number  # 检索出8张相似度最高的图片
    img_list = []
    for i, index in enumerate(rank_ID[0:maxres]):
        img_list.append(imgNames[index])
    img_list=[str(i,'utf-8') for i in img_list]
    return img_list


if __name__ == "__main__":
    print(get_img_list('img/1.jpg',120))